# This file is part of Cantera. See License.txt in the top-level directory or
# at https://cantera.org/license.txt for license and copyright information.

from typing import Any, Literal, TypeAlias, TypedDict, overload

from typing_extensions import Required

from ._types import (
    Array,
    ArrayLike,
    Basis,
    CompositionLike,
    CompositionVariable,
    EquilibriumSolver,
    FullState,
    LogLevel,
    PropertyPair,
    State2Setter,
    StateSetter,
    StateVariable,
)
from .composite import Solution
from .solutionbase import _SolutionBase
from .speciesthermo import SpeciesThermo, _SpeciesThermoInput
from .transport import GasTransportData, _GasTransportInput
from .units import Units

# Avoid fixed options unless we can find a way to support custom extensions:
# _ThermoType: TypeAlias = Literal[
#     "Debye-Huckel",
#     "HMW-electrolyte",
#     "Margules",
#     "MixtureFugacity",
#     "Peng-Robinson",
#     "Phase",
#     "Redlich-Kister",
#     "Redlich-Kwong",
#     "SingleSpecies",
#     "none",
#     "binary-solution-tabulated",
#     "coverage-dependent-surface",
#     "edge",
#     "electron-cloud",
#     "fixed-stoichiometry",
#     "ideal-condensed",
#     "ideal-gas",
#     "ideal-molal-solution",
#     "ideal-solution-VPSS",
#     "ideal-surface",
#     "liquid-water-IAPWS95",
#     "plasma",
#     "pure-fluid",
# ]
_ThermoType: TypeAlias = str

# _PhaseOfMatter: TypeAlias = Literal[
#     "gas",
#     "liquid",
#     "solid",
#     "supercritical",
#     "unstable-liquid",
#     "unstable-gas",
#     "liquid-gas-mix",
#     "unspecified",
# ]
_PhaseOfMatter: TypeAlias = str

_QuadratureMethod: TypeAlias = Literal["simpson", "trapezoidal"]

_SpeciesInput = TypedDict(
    "_SpeciesInput",
    {
        "name": Required[str],
        "composition": Required[dict[str, float]],
        "thermo": _SpeciesThermoInput,
        "transport": _GasTransportInput,
        "equation-of-state": dict[str, Any],
        "critical-parameters": dict[str, float],
        "Debye-Huckel": dict[str, float],
    },
    total=False,
)

class ThermoModelMethodError(Exception):
    def __init__(self, thermo_model: str) -> None: ...

class Species:
    def __init__(
        self,
        name: str | None = None,
        composition: CompositionLike | None = None,
        charge: float | None = None,
        size: float | None = None,
        init: bool = True,
    ) -> None: ...
    @staticmethod
    def from_dict(data: _SpeciesInput) -> Species: ...
    @staticmethod
    def from_yaml(text: str) -> Species: ...
    @staticmethod
    def list_from_file(filename: str, section: str = "species") -> list[Species]: ...
    @staticmethod
    def list_from_yaml(text: str, section: str | None = None) -> list[Species]: ...
    @property
    def name(self) -> str: ...
    @property
    def composition(self) -> dict[str, float]: ...
    @property
    def charge(self) -> float: ...
    @property
    def size(self) -> float: ...
    @property
    def molecular_weight(self) -> float: ...
    @property
    def thermo(self) -> SpeciesThermo: ...
    @thermo.setter
    def thermo(self, value: SpeciesThermo) -> None: ...
    @property
    def transport(self) -> GasTransportData: ...
    @transport.setter
    def transport(self, value: GasTransportData) -> None: ...
    @property
    def input_data(self) -> _SpeciesInput: ...
    def update_user_data(self, data: dict[str, Any]) -> None: ...
    def clear_user_data(self) -> None: ...

class ThermoPhase(_SolutionBase):
    @property
    def thermo_model(self) -> _ThermoType: ...
    @property
    def phase_of_matter(self) -> _PhaseOfMatter: ...
    def report(self, show_thermo: bool = True, threshold: float = 1e-14) -> str: ...
    def __call__(self, *args: Any, **kwargs: Any) -> None: ...
    @property
    def is_pure(self) -> bool: ...
    @property
    def has_phase_transition(self) -> bool: ...
    @property
    def is_compressible(self) -> bool: ...
    @property
    def _native_mode(self) -> FullState: ...
    @property
    def _native_state(
        self,
    ) -> tuple[StateVariable, StateVariable, CompositionVariable]: ...
    @property
    def _full_states(
        self,
    ) -> dict[frozenset[StateVariable | CompositionVariable], FullState]: ...
    @property
    def _partial_states(self) -> dict[frozenset[StateVariable], PropertyPair]: ...
    @property
    def basis(self) -> Basis: ...
    @basis.setter
    def basis(self, value: Basis) -> None: ...
    def equilibrate(
        self,
        XY: PropertyPair,
        solver: EquilibriumSolver = "auto",
        rtol: float = 1e-9,
        max_steps: int = 1000,
        max_iter: int = 100,
        estimate_equil: int = 0,
        log_level: LogLevel = 0,
    ) -> None: ...
    ####### Composition, species, and elements ########
    @property
    def n_elements(self) -> int: ...
    def element_index(self, element: str | int) -> int: ...
    def element_name(self, m: int) -> str: ...
    @property
    def element_names(self) -> list[str]: ...
    def atomic_weight(self, m: str | int) -> float: ...
    @property
    def atomic_weights(self) -> Array: ...
    @property
    def n_species(self) -> int: ...
    @property
    def n_selected_species(self) -> int: ...
    def species_name(self, k: int) -> str: ...
    @property
    def species_names(self) -> list[str]: ...
    def species_index(self, species: str | int) -> int: ...
    @property
    def case_sensitive_species_names(self) -> bool: ...
    @case_sensitive_species_names.setter
    def case_sensitive_species_names(self, val: bool) -> None: ...
    @overload
    def species(self, k: None = None) -> list[Species]: ...
    @overload
    def species(self, k: str | int) -> Species: ...
    @overload
    def species(self, k: str | int | None = None) -> Species | list[Species]: ...
    def modify_species(self, k: int, species: Species) -> None: ...
    def add_species(self, species: Species) -> None: ...
    def add_species_alias(self, name: str, alias: str) -> None: ...
    def find_isomers(self, comp: CompositionLike) -> list[str]: ...
    def n_atoms(self, species: str | int, element: str | int) -> int: ...
    @property
    def molecular_weights(self) -> Array: ...
    @property
    def charges(self) -> Array: ...
    @property
    def mean_molecular_weight(self) -> float: ...
    @property
    def Y(self) -> Array: ...
    @Y.setter
    def Y(self, Y: CompositionLike) -> None: ...
    @property
    def X(self) -> Array: ...
    @X.setter
    def X(self, X: CompositionLike) -> None: ...
    @property
    def concentrations(self) -> Array: ...
    @concentrations.setter
    def concentrations(self, C: ArrayLike) -> None: ...
    def __composition_to_array(self, comp: CompositionLike, basis: Basis) -> Array: ...
    def set_equivalence_ratio(
        self,
        phi: float,
        fuel: CompositionLike,
        oxidizer: CompositionLike,
        basis: Basis = "mole",
        *,
        diluent: CompositionLike | None = None,
        fraction: str
        | dict[Literal["fuel", "oxidizer", "diluent"], float]
        | None = None,
    ) -> None: ...
    def set_mixture_fraction(
        self,
        mixture_fraction: float,
        fuel: CompositionLike,
        oxidizer: CompositionLike,
        basis: Basis = "mole",
    ) -> None: ...
    def equivalence_ratio(
        self,
        fuel: CompositionLike | None = None,
        oxidizer: CompositionLike | None = None,
        basis: Basis = "mole",
        include_species: list[str | int] | None = None,
    ) -> float: ...
    def mixture_fraction(
        self,
        fuel: CompositionLike,
        oxidizer: CompositionLike,
        basis: Basis = "mole",
        element: str | int = "Bilger",
    ) -> float: ...
    def stoich_air_fuel_ratio(
        self,
        fuel: CompositionLike,
        oxidizer: CompositionLike,
        basis: Basis = "mole",
    ) -> float: ...
    def elemental_mass_fraction(self, m: str | int) -> float: ...
    def elemental_mole_fraction(self, m: str | int) -> float: ...
    def set_unnormalized_mass_fractions(self, Y: ArrayLike) -> None: ...
    def set_unnormalized_mole_fractions(self, X: ArrayLike) -> None: ...
    def mass_fraction_dict(self, threshold: float = 0.0) -> dict[str, float]: ...
    def mole_fraction_dict(self, threshold: float = 0.0) -> dict[str, float]: ...
    ######## Read-only thermodynamic properties ########
    @property
    def P(self) -> float: ...
    @property
    def T(self) -> float: ...
    @property
    def density(self) -> float: ...
    @property
    def density_mass(self) -> float: ...
    @property
    def density_mole(self) -> float: ...
    @property
    def v(self) -> float: ...
    @property
    def volume_mass(self) -> float: ...
    @property
    def volume_mole(self) -> float: ...
    @property
    def u(self) -> float: ...
    @property
    def int_energy_mole(self) -> float: ...
    @property
    def int_energy_mass(self) -> float: ...
    @property
    def h(self) -> float: ...
    @property
    def enthalpy_mole(self) -> float: ...
    @property
    def enthalpy_mass(self) -> float: ...
    @property
    def s(self) -> float: ...
    @property
    def entropy_mole(self) -> float: ...
    @property
    def entropy_mass(self) -> float: ...
    @property
    def g(self) -> float: ...
    @property
    def gibbs_mole(self) -> float: ...
    @property
    def gibbs_mass(self) -> float: ...
    @property
    def cv(self) -> float: ...
    @property
    def cv_mole(self) -> float: ...
    @property
    def cv_mass(self) -> float: ...
    @property
    def cp(self) -> float: ...
    @property
    def cp_mole(self) -> float: ...
    @property
    def cp_mass(self) -> float: ...
    @property
    def critical_temperature(self) -> float: ...
    @property
    def critical_pressure(self) -> float: ...
    @property
    def critical_density(self) -> float: ...
    @property
    def P_sat(self) -> float: ...
    @property
    def T_sat(self) -> float: ...
    @property
    def auxiliary_data(self) -> dict[str, Any]: ...
    ######## Methods to get/set the complete thermodynamic state ########
    @property
    def state_size(self) -> int: ...
    @property
    def state(self) -> Array: ...
    @state.setter
    def state(self, state: ArrayLike) -> None: ...
    @property
    def TD(self) -> tuple[float, float]: ...
    @TD.setter
    def TD(self, values: State2Setter) -> None: ...
    @property
    def TDX(self) -> tuple[float, float, Array]: ...
    @TDX.setter
    def TDX(self, values: StateSetter) -> None: ...
    @property
    def TDY(self) -> tuple[float, float, Array]: ...
    @TDY.setter
    def TDY(self, values: StateSetter) -> None: ...
    @property
    def TP(self) -> tuple[float, float]: ...
    @TP.setter
    def TP(self, values: State2Setter) -> None: ...
    @property
    def TPX(self) -> tuple[float, float, Array]: ...
    @TPX.setter
    def TPX(self, values: StateSetter) -> None: ...
    @property
    def TPY(self) -> tuple[float, float, Array]: ...
    @TPY.setter
    def TPY(self, values: StateSetter) -> None: ...
    @property
    def UV(self) -> tuple[float, float]: ...
    @UV.setter
    def UV(self, values: State2Setter) -> None: ...
    @property
    def UVX(self) -> tuple[float, float, Array]: ...
    @UVX.setter
    def UVX(self, values: StateSetter) -> None: ...
    @property
    def UVY(self) -> tuple[float, float, Array]: ...
    @UVY.setter
    def UVY(self, values: StateSetter) -> None: ...
    @property
    def DP(self) -> tuple[float, float]: ...
    @DP.setter
    def DP(self, values: State2Setter) -> None: ...
    @property
    def DPX(self) -> tuple[float, float, Array]: ...
    @DPX.setter
    def DPX(self, values: StateSetter) -> None: ...
    @property
    def DPY(self) -> tuple[float, float, Array]: ...
    @DPY.setter
    def DPY(self, values: StateSetter) -> None: ...
    @property
    def HP(self) -> tuple[float, float]: ...
    @HP.setter
    def HP(self, values: State2Setter) -> None: ...
    @property
    def HPX(self) -> tuple[float, float, Array]: ...
    @HPX.setter
    def HPX(self, values: StateSetter) -> None: ...
    @property
    def HPY(self) -> tuple[float, float, Array]: ...
    @HPY.setter
    def HPY(self, values: StateSetter) -> None: ...
    @property
    def SP(self) -> tuple[float, float]: ...
    @SP.setter
    def SP(self, values: State2Setter) -> None: ...
    @property
    def SPX(self) -> tuple[float, float, Array]: ...
    @SPX.setter
    def SPX(self, values: StateSetter) -> None: ...
    @property
    def SPY(self) -> tuple[float, float, Array]: ...
    @SPY.setter
    def SPY(self, values: StateSetter) -> None: ...
    @property
    def SV(self) -> tuple[float, float]: ...
    @SV.setter
    def SV(self, values: State2Setter) -> None: ...
    @property
    def SVX(self) -> tuple[float, float, Array]: ...
    @SVX.setter
    def SVX(self, values: StateSetter) -> None: ...
    @property
    def SVY(self) -> tuple[float, float, Array]: ...
    @SVY.setter
    def SVY(self, values: StateSetter) -> None: ...
    # partial molar / non-dimensional properties
    @property
    def partial_molar_enthalpies(self) -> Array: ...
    @property
    def partial_molar_entropies(self) -> Array: ...
    @property
    def partial_molar_int_energies(self) -> Array: ...
    @property
    def chemical_potentials(self) -> Array: ...
    @property
    def electrochemical_potentials(self) -> Array: ...
    @property
    def partial_molar_cp(self) -> Array: ...
    @property
    def partial_molar_volumes(self) -> Array: ...
    @property
    def standard_enthalpies_RT(self) -> Array: ...
    @property
    def standard_entropies_R(self) -> Array: ...
    @property
    def standard_int_energies_RT(self) -> Array: ...
    @property
    def standard_gibbs_RT(self) -> Array: ...
    @property
    def standard_cp_R(self) -> Array: ...
    @property
    def activities(self) -> Array: ...
    @property
    def activity_coefficients(self) -> Array: ...
    ######## Miscellaneous properties ########
    @property
    def isothermal_compressibility(self) -> float: ...
    @property
    def thermal_expansion_coeff(self) -> float: ...
    @property
    def sound_speed(self) -> float: ...
    @property
    def min_temp(self) -> float: ...
    @property
    def max_temp(self) -> float: ...
    @property
    def reference_pressure(self) -> float: ...
    @property
    def electric_potential(self) -> float: ...
    @electric_potential.setter
    def electric_potential(self, value: float) -> None: ...
    @property
    def standard_concentration_units(self) -> Units: ...
    # methods for plasma
    @property
    def Te(self) -> float: ...
    @Te.setter
    def Te(self, value: float) -> None: ...
    @property
    def Pe(self) -> float: ...
    @property
    def reduced_electric_field(self) -> float: ...
    @reduced_electric_field.setter
    def reduced_electric_field(self, value: float) -> None: ...
    @property
    def electric_field(self) -> float: ...
    @electric_field.setter
    def electric_field(self, value: float) -> None: ...
    def set_discretized_electron_energy_distribution(
        self, levels: ArrayLike, distribution: ArrayLike
    ) -> None: ...
    def update_electron_energy_distribution(self) -> None: ...
    @property
    def n_electron_energy_levels(self) -> int: ...
    @property
    def electron_energy_levels(self) -> Array: ...
    @electron_energy_levels.setter
    def electron_energy_levels(self, levels: Array) -> None: ...
    @property
    def electron_energy_distribution(self) -> Array: ...
    @property
    def isotropic_shape_factor(self) -> float: ...
    @isotropic_shape_factor.setter
    def isotropic_shape_factor(self, x: float) -> None: ...
    @property
    def electron_energy_distribution_type(self) -> str: ...
    @electron_energy_distribution_type.setter
    def electron_energy_distribution_type(self, distribution_type: str) -> None: ...
    @property
    def mean_electron_energy(self) -> float: ...
    @mean_electron_energy.setter
    def mean_electron_energy(self, energy: float) -> None: ...
    @property
    def quadrature_method(self) -> _QuadratureMethod: ...
    @quadrature_method.setter
    def quadrature_method(self, method: _QuadratureMethod) -> None: ...
    @property
    def normalize_electron_energy_distribution_enabled(self) -> bool: ...
    @normalize_electron_energy_distribution_enabled.setter
    def normalize_electron_energy_distribution_enabled(self, enable: bool) -> None: ...
    @property
    def electron_species_name(self) -> str: ...
    @property
    def elastic_power_loss(self) -> float: ...

class InterfacePhase(ThermoPhase):
    @property
    def adjacent(self) -> dict[str, Solution]: ...
    @property
    def site_density(self) -> float: ...
    @site_density.setter
    def site_density(self, value: float) -> None: ...
    @property
    def coverages(self) -> Array: ...
    @coverages.setter
    def coverages(self, theta: CompositionLike) -> None: ...
    def set_unnormalized_coverages(self, cov: ArrayLike) -> None: ...

class PureFluid(ThermoPhase):
    @property
    def Q(self) -> float: ...
    @Q.setter
    def Q(self, Q: float) -> None: ...
    @property
    def TQ(self) -> tuple[float, float]: ...
    @TQ.setter
    def TQ(self, values: State2Setter) -> None: ...
    @property
    def PQ(self) -> tuple[float, float]: ...
    @PQ.setter
    def PQ(self, values: State2Setter) -> None: ...
    @property
    def ST(self) -> tuple[float, float]: ...
    @ST.setter
    def ST(self, values: State2Setter) -> None: ...
    @property
    def TV(self) -> tuple[float, float]: ...
    @TV.setter
    def TV(self, values: State2Setter) -> None: ...
    @property
    def PV(self) -> tuple[float, float]: ...
    @PV.setter
    def PV(self, values: State2Setter) -> None: ...
    @property
    def UP(self) -> tuple[float, float]: ...
    @UP.setter
    def UP(self, values: State2Setter) -> None: ...
    @property
    def VH(self) -> tuple[float, float]: ...
    @VH.setter
    def VH(self, values: State2Setter) -> None: ...
    @property
    def TH(self) -> tuple[float, float]: ...
    @TH.setter
    def TH(self, values: State2Setter) -> None: ...
    @property
    def SH(self) -> tuple[float, float]: ...
    @SH.setter
    def SH(self, values: State2Setter) -> None: ...
    @property
    def TDQ(self) -> tuple[float, float, float]: ...
    @property
    def TPQ(self) -> tuple[float, float, float]: ...
    @TPQ.setter
    def TPQ(self, values: tuple[float | None, float | None, float | None]) -> None: ...
    @property
    def UVQ(self) -> tuple[float, float, float]: ...
    @property
    def DPQ(self) -> tuple[float, float, float]: ...
    @property
    def HPQ(self) -> tuple[float, float, float]: ...
    @property
    def SPQ(self) -> tuple[float, float, float]: ...
    @property
    def SVQ(self) -> tuple[float, float, float]: ...

class Element:
    num_elements_defined: int
    element_symbols: tuple[str, ...]
    element_names: tuple[str, ...]
    def __init__(self, arg: str | int) -> None: ...
    @property
    def name(self) -> str: ...
    @property
    def atomic_number(self) -> int: ...
    @property
    def symbol(self) -> str: ...
    @property
    def weight(self) -> float: ...
